﻿using AiBLSmartEdu.Module.AIMedicalAssistant.API.Configurations;
using ChromaDB.Client;
using FrameworkCore.DI;
using Microsoft.Extensions.Options;

#nullable disable

namespace AiBLSmartEdu.Module.AIMedicalAssistant.API.Services;

public class ChromaCollectionService : IDisposable, ITransientDependency
{
    private readonly HttpClient _httpClient;
    private readonly string _collectionName;
    private ChromaCollectionClient _client;
    private IOptions<ChromaSettings> _options;

    public ChromaCollectionService(
        IOptions<ChromaSettings> options,
        IHttpClientFactory httpClientFactory)
    {
        _collectionName = options.Value.CollectionName;
        _httpClient = httpClientFactory.CreateClient("Chroma");
        InitializeClient().Wait();
    }

    private async Task InitializeClient()
    {
        var config = new ChromaConfigurationOptions(uri: new Uri(_options.Value.ApiUrl));
        var client = new ChromaClient(config, _httpClient);
        var collection = await client.GetOrCreateCollection(_collectionName);
        _client = new ChromaCollectionClient(collection, config, _httpClient);
    }

    public async Task BulkUpsertAsync(
        IEnumerable<string> ids,
        IEnumerable<float[]> embeddings,
        IEnumerable<Dictionary<string, object>> metadatas)
    {
        await _client.Upsert(
            ids: ids.ToList(),
            embeddings: embeddings.Select(e => new ReadOnlyMemory<float>(e)).ToList(),
            metadatas: metadatas.ToList()
        );
    }

    public void Dispose() => _httpClient.Dispose();
}